# -*- coding: utf-8 -*-
# @Author  : longbhu
# @Time    : 2025/3/21 13:40
# @Function: 第一步：将2010land30m合并成31个省；将每个省替换代码，替换代码为新国标
"""
E:\land30m2010\anhui

"""

import os
import numpy as np
import rioxarray as rxr
import rasterio
from rasterio.merge import merge


def merge_tifs_in_directory(province_dir: str, output_file: str) -> None:
    """
    合并指定目录中的所有 TIFF 文件，并将结果写入到输出文件中。

    参数:
    province_dir (str): 包含 TIFF 文件的目录路径。
    output_file (str): 合并结果的输出文件路径。
    """
    # 使用 rasterio.open 读取 TIFF 文件
    rasters = [rasterio.open(os.path.join(province_dir, f)) for f in os.listdir(province_dir) if f.endswith('.tif')]

    if not rasters:
        print("没有找到有效的 TIFF 文件")
        return

    # 合并 TIFF 文件
    merged_raster, out_trans = merge(rasters)

    # 写入合并结果到输出文件
    with rasterio.open(
            output_file,
            'w',
            driver='GTiff',
            height=merged_raster.shape[1],
            width=merged_raster.shape[2],
            count=merged_raster.shape[0],
            dtype=np.uint8,  # 设置数据类型为 uint8
            crs=rasters[0].crs,
            transform=out_trans,
            nodata=255,
            # compress='lzw'
    ) as dst:
        dst.write(merged_raster)

# Example usage
base_dir = r'E:\land30m2010'
for province in os.listdir(base_dir):
    province_dir = os.path.join(base_dir, province)
    if os.path.isdir(province_dir):
        base_dir_h = r'H:\land30m2010'
        output_file = os.path.join(base_dir_h, f'{province}_merged.tif')
        print(f'Merging files in {province_dir} to {output_file}')
        merge_tifs_in_directory(province_dir, output_file)


